import argparse
import logging
import os
from typing import List, Optional

import dashscope
from dashscope.audio.tts_v2 import SpeechSynthesizer, ResultCallback


class QwenSynthesisCallback(ResultCallback):
    def __init__(self, output_file):
        self.output_file = output_file
        self.file = None
        
    def on_open(self):
        self.file = open(self.output_file, "wb")
        
    def on_complete(self):
        pass
        
    def on_error(self, message: str):
        logging.error(f"Speech synthesis task failed: {message}")
        
    def on_close(self):
        if self.file:
            self.file.close()
            
    def on_event(self, message):
        pass
        
    def on_data(self, data: bytes) -> None:
        if self.file:
            self.file.write(data)


class QwenTTSProvider:
    """Provider for Qwen text-to-speech synthesis."""

    def __init__(self, api_key: Optional[str] = None):
        """Initialize the Qwen TTS provider.

        Args:
            api_key: The Dashscope API key. If not provided, will try to get from environment.
        """
        self.api_key = api_key
        if self.api_key:
            dashscope.api_key = self.api_key
            
        self.model = "cosyvoice-v2"
        self.voice = "longxiaochun_v2"

    def synthesize_text(
        self,
        text: str,
        output_file: str,
        voice: Optional[str] = None,
        audio_format: str = "mp3",
    ) -> str:
        """Synthesize a single text utterance to audio file.

        Args:
            text: The text to synthesize
            output_file: Path to save the audio file
            voice: Optional voice to use (overrides default)
            audio_format: Audio format (mp3, wav, etc.)

        Returns:
            Path to the generated audio file
        """
        # Ensure output directory exists
        os.makedirs(os.path.dirname(os.path.abspath(output_file)), exist_ok=True)
        
        # Create callback for writing to file
        callback = QwenSynthesisCallback(output_file)
        
        # Set voice if provided, otherwise use default
        selected_voice = voice if voice else self.voice
        
        # Create synthesizer
        synthesizer = SpeechSynthesizer(
            model=self.model,
            voice=selected_voice,
            callback=callback,
        )
        
        # Call the synthesizer
        synthesizer.call(text)
        
        return output_file

    def synthesize_from_file(
        self,
        input_file: str,
        output_dir: str,
        voice: Optional[str] = None,
        audio_format: str = "mp3",
    ) -> List[str]:
        """Synthesize audio for each line in a text file.

        Args:
            input_file: Path to the input text file (one utterance per line)
            output_dir: Directory to save the generated audio files
            voice: Optional voice to use
            audio_format: Audio format (mp3, wav, etc.)

        Returns:
            List of paths to the generated audio files
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        output_files = []

        with open(input_file, "r", encoding="utf-8") as f:
            lines = [line.strip() for line in f.readlines() if line.strip()]

        for i, line in enumerate(lines):
            output_file = os.path.join(output_dir, f"audio_{i + 1}.{audio_format}")
            self.synthesize_text(line, output_file, voice, audio_format)
            output_files.append(output_file)

        return output_files


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Qwen TTS synthesizer")
    parser.add_argument(
        "input_file", help="Path to input text file (one utterance per line)"
    )
    parser.add_argument(
        "output_dir", help="Directory to save the generated audio files"
    )
    parser.add_argument(
        "--voice", help="Optional voice to use", default=None
    )
    parser.add_argument("--format", help="Audio format (mp3, wav, etc.)", default="mp3")
    parser.add_argument(
        "--api_key", help="Dashscope API key (if not set in environment)", default=None
    )

    args = parser.parse_args()

    provider = QwenTTSProvider(api_key=args.api_key)
    output_files = provider.synthesize_from_file(
        input_file=args.input_file,
        output_dir=args.output_dir,
        voice=args.voice,
        audio_format=args.format,
    )

    print(f"Generated {len(output_files)} audio files in {args.output_dir}")
    for file in output_files:
        print(f"- {file}")
